#include "solver.h"


// BANDWIDTH->(BLOCK_SIZE, TILE_SIZE)
// 论文中折线图的实现
void set_size(MATRIX *A, int block_size, int tile_size)
{

    A->bandwidth = get_bw(A, A->real_n);
    printf("bandwidth = %d\n", A->bandwidth);
    printf("expected block_size = %.1lf, tile_size = %.1lf\n", 0.08*A->bandwidth, 0.02*A->bandwidth);
    if (block_size > 0 && tile_size > 0)
    {
        A->block_size = block_size;
        A->tile_size = tile_size;
    }
    else if (A->bandwidth <= 256)
    {
        A->block_size = 1;
        A->tile_size = 48;
    }
    else if (A->bandwidth <= 1024)
    {
        A->block_size = 2;
        A->tile_size = 64;
    }
    else if (A->bandwidth <= 384 * 8)
    {
        A->block_size = 2;
        A->tile_size = 96;
    }
    else if (A->bandwidth <= 384 * 16)
    {
        A->block_size = 3;
        A->tile_size = 128;
    }
    else
    {
        A->block_size = 2;
        A->tile_size = 192;
    }
}

// 设置val为当前时间
double set_time(struct timeval *val)
{
    gettimeofday(val, NULL);
    return 0.0;
}

// 计算start到现在的时间间隔(s)
double get_time(struct timeval *start)
{
    struct timeval now;
    gettimeofday(&now, NULL);
    double time = now.tv_sec - start->tv_sec + (now.tv_usec - start->tv_usec) * 1e-6;
    return time;
}

void sort(MATRIX *A)
{
    if (!A || A->type != CSR)
        return;
    for (int i = 0; i < A->real_m; i++)
    {
        for (int j = A->row[i]; j < A->row[i + 1]; j++)
        {
            int key = A->col[j];
            double val = A->a[j];
            int k = j - 1;
            while (k >= A->row[i] && A->col[k] > key)
            {
                A->col[k + 1] = A->col[k];
                A->a[k + 1] = A->a[k];
                k--;
            }
            A->col[k + 1] = key;
            A->a[k + 1] = val;
        }
    }
}

void permAperm_inv(MATRIX *A, int *perm)
{
    if (!A || A->type != CSR)
        return;

    int *perm_inv = (int *)my_malloc(A->real_m, sizeof(int));
    for (int i = 0; i < A->real_m; ++i)
        perm_inv[perm[i]] = i;

    int *row = (int *)my_malloc(A->real_m + 1, sizeof(int));
    int *col = (int *)my_malloc(A->real_nonz, sizeof(int));
    double *a = (double *)my_malloc(A->real_nonz, sizeof(double));

    // perm, row exchange
    row[0] = 0;
    for (int i = 0; i < A->real_m; i++)
    {
        row[i + 1] = A->row[perm[i] + 1] - A->row[perm[i]];
    }
    for (int i = 0; i < A->real_m; i++)
    {
        row[i + 1] += row[i];
    }

    // perm_inv, column exchange
    for (int i = 0; i < A->real_m; i++)
    {
        int j1 = A->row[perm[i]];
        int j2 = row[i];
        while (j2 < row[i + 1])
        {
            col[j2] = perm_inv[A->col[j1]];
            a[j2] = A->a[j1];
            j1++;
            j2++;
        }
    }

    my_free(perm_inv, A->real_m, sizeof(int));
    my_free(A->row, A->real_m + 1, sizeof(int));
    my_free(A->col, A->real_nonz, sizeof(int));
    my_free(A->a, A->real_nonz, sizeof(double));

    A->row = row;
    A->col = col;
    A->a = a;
    sort(A);
}

void permTILE(MATRIX *A, int *perm)
{
    if (!A || A->type == NONE || A->type == CSR)
        return;

    int inew, jnew, binew, bjnew, tinew, tjnew;
    BLOCK **block = (BLOCK **)my_malloc(A->m * A->n, sizeof(BLOCK *));

    BLOCK *B1, *B2;
    for (int bi = 0; bi < A->m; bi++)
    {
        for (int bj = 0; bj < A->n; bj++)
        {
            int m = get_bs(A, bi);
            int n = get_bs(A, bj);
            for (int ti = 0; ti < m; ti++)
            {
                for (int tj = 0; tj < n; tj++)
                {

                    inew = perm[bi * A->block_size + ti];
                    binew = inew / A->block_size;
                    tinew = inew % A->block_size;

                    jnew = perm[bj * A->block_size + tj];
                    bjnew = jnew / A->block_size;
                    tjnew = jnew % A->block_size;

                    B1 = A->block[binew * A->n + bjnew];
                    if (!B1)
                        continue;

                    B2 = block[bi * A->n + bj];
                    if (!B2)
                    {
                        B2 = block[bi * A->n + bj] = construct_B(m, n);
                    }
                    B2->tile[ti * B2->n + tj] = B1->tile[tinew * B1->n + tjnew];
                }
            }
        }
    }

    for (int bi = 0; bi < A->m; bi++)
    {
        for (int bj = 0; bj < A->n; bj++)
        {
            my_free(A->block[bi * A->n + bj], 1, sizeof(BLOCK));
        }
    }

    A->block = block;
}

void permA(MATRIX *A, int *perm)
{
    if (!A || A->type == NONE)
        return;

    if (A->type == CSR)
    {
        permAperm_inv(A, perm);
    }
    else
    {
        permTILE(A, perm);
    }
}

// get_block_size, 获取第b行/列BLOCK对应的tile行/列数
// 最后一行/列BLOCK可能不完整，其他BLOCK均返回A->block_size
int get_bs(const MATRIX *A, int b)
{
    int ret = A->real_n - A->block_size * A->tile_size * b;
    ret = MIN(ret, A->block_size * A->tile_size);

    int bs = ret / A->tile_size;
    if (ret % A->tile_size)
        bs++;

    // MIN似乎是非必要的
    return MIN(bs, A->block_size);
}

// get_block_size, 获取第b行/列BLOCK内第t行/列TILE对应的元素行/列数
// 最后一行/列TILE可能不完整，其他TILE均返回A->tile_size
int get_ts(const MATRIX *A, int b, int t)
{
    int ret = A->real_n - A->block_size * A->tile_size * b - A->tile_size * t;
    return MIN(ret, A->tile_size);
}

// 将全局坐标i转换为三元坐标(b, t, r)
// b=block_index, t=tile_index, r=inner_tile_index
void get_position(const MATRIX *A, int i, int *b, int *t, int *r)
{
    int epb = A->block_size * A->tile_size;
    int ept = A->tile_size;
    *b = i / epb;
    *t = (i % epb) / ept;
    *r = (i % epb) % ept;
}

void print_row(MATRIX *A, int bi, int ti, int ai)
{
    BLOCK *B;
    TILE *T;
    int bj, tj, aj;
    for (int j = 0; j < A->real_n; j++)
    {
        get_position(A, j, &bj, &tj, &aj);
        B = A->block[bi * A->n + bj];
        if (!B)
            continue;
        T = B->tile[ti * B->n + tj];
        if (!T)
            continue;
        double x = T->a[ai * T->n + aj];
        if (x == 0)
            continue;
        printf("%e ", x);
    }
    printf("\n");
}
void print_col(MATRIX *A, int bj, int tj, int aj)
{
    BLOCK *B;
    TILE *T;
    int bi, ti, ai;
    for (int i = 0; i < A->real_m; i++)
    {
        get_position(A, i, &bi, &ti, &ai);
        B = A->block[bi * A->n + bj];
        if (!B)
            continue;
        T = B->tile[ti * B->n + tj];
        if (!T)
            continue;
        double x = T->a[ai * T->n + aj];
        if (x == 0)
            continue;
        printf("%e ", x);
    }
    printf("\n");
}
void print_TILE(TILE *T)
{
    if (!T)
    {
        printf("NULL TILE\n");
        return;
    }
    const char *blank = "         ";

    for (int i = 0; i < T->m; i++)
    {
        printf("%d: ", i % 10);
        for (int j = 0; j < T->n; j++)
        {
            if (T->a[i * T->n + j] == 0)
            {
                printf("%s", blank);
            }
            else if (T->a[i * T->n + j] > 0)
            {
                printf(" %.1e ", T->a[i * T->n + j]);
            }
            else
            {
                printf("%.1e ", T->a[i * T->n + j]);
            }
        }
        printf("\n");
    }
}

void print_BLOCK(BLOCK *B)
{
    if (!B)
    {
        printf("NULL BLOCK\n");
        return;
    }
    for (int i = 0; i < B->m; i++)
    {
        for (int j = 0; j < B->n; j++)
        {
            if (!B->tile[i * B->n + j])
                continue;
            printf("TILE[%d, %d]:\n", i, j);
            print_TILE(B->tile[i * B->n + j]);
        }
    }
}

void print_MATRIX(MATRIX *A)
{
    if (!A)
    {
        printf("NULL MATRIX\n");
        return;
    }
    for (int i = 0; i < A->m; i++)
    {
        for (int j = 0; j < A->n; j++)
        {
            if (!A->block[i * A->n + j])
                continue;
            printf("BLOCK[%d, %d]:\n", i, j);
            print_BLOCK(A->block[i * A->n + j]);
        }
    }
}

void progress_bar(const char *name, int total, int n)
{
    int length = 50;
    char icon = '-';

    total--;
    int count = ((double)n / total) * length;

    printf("%s: [", name); // 3 + strlen(name)
    for (int i = 0; i < length; i++)
    {
        if (i < count)
        {
            printf("%c", icon);
        }
        else
        {
            printf(" ");
        }
    }
    printf("] "); // length 2

    printf("%3d%%", count * 2); // length 4
    fflush(stdout);

    if (n < total)
    {
        for (int i = 0; i < 3 + strlen(name) + length + 2 + 4; i++)
        {
            printf("\b");
        }
        fflush(stdout);
    }
    else
    {
        printf("\n");
    }
}

// 计算TILE中非零元素数量
void update_nonz(TILE *T)
{
    T->nonz = 0;
    for (int i = 0; i < T->m; i++)
    {
        for (int j = 0; j < T->n; j++)
        {
            if (T->a[i * T->n + j] == 0)
                continue;
            T->nonz++;
        }
    }
}

int scatter(const MATRIX *A)
{
    if (!A || A->type == NONE)
        return 0;

    const char *py = "python3 ";
    const char *script = "scatter.py ";

    char file[64];
    if (A->type == CSR)
    {
        sprintf(file, "./cache/pattern/%s.txt", A->name);
    }
    else
    {
        sprintf(file, "./cache/pattern/block_%s.txt", A->name);
    }

    char command[64];
    strcpy(command, py);
    strcat(command, script);
    strcat(command, file);

    return system(command);
}

void pattern2(const MATRIX* A){
    if(!A || A->type != CSC) {
        printf("pattern2: A is NULL or A is not CSC\n");
        return;
    }

    double scale = MIN((double)1000.0 / A->real_n, (double)1.0);
    // double scale = 1.0;
    if (A->real_n <= 0 || A->row == NULL || A->col == NULL || scale <= 0) return;

    std::set< std::pair<int, int> > coo;

    for (int j = 0; j < A->real_n; j++) {
        for (int i = A->col[j]; i < A->col[j + 1]; i++) {
            int x = A->row[i] * scale;
            int y = j * scale;
            coo.insert( std::make_pair(x, y) );
        }
    }

    fs::path dir = "./cache/pattern";
    fs::path filepath = dir / "data.txt";
    if (!fs::exists(dir)) {
        fs::create_directories(dir);
    }
    std::ofstream ofs(filepath);
    if (!ofs.is_open()) {
        printf("open \"%s\" failed, please check if the directory exists.\n", filepath.string().c_str());
        return;
    }
    for (const auto &c : coo) {
        ofs << c.first << " " << c.second << " ";
    }
    ofs.close();

    // 将python3 pattern.py 和 output_filename拼接起来
    char cmd[512];
    sprintf(cmd, "python3 pattern.py %s", A->name);
    if(system(cmd) == -1) {
        printf("%s, pattern fail\n", cmd);
    }
}

void pattern(const MATRIX *A)
{
    if (!A || A->type == NONE)
        return;

    char file[64];
    if (A->type == CSR)
    {
        sprintf(file, "./cache/pattern/%s.txt", A->name);
    }
    else
    {
        sprintf(file, "./cache/pattern/block_%s.txt", A->name);
    }

    FILE *fp = fopen(file, "w");
    if (!fp)
        ABORT("fopen fail");

    fprintf(fp, "%d %d %d\n", A->real_m, A->real_n, A->real_nonz);
    switch (A->type)
    {
    case CSR:
        for (int i = 0; i < A->real_m; i++)
        {
            for (int j = A->row[i]; j < A->row[i + 1]; j++)
            {
#if 0
					fprintf(fp, "%d %d %f\n", i, A->col[j], A->a[j]);
#else
                fprintf(fp, "%d %d %f\n", i, A->col[j], 1.0);
#endif
            }
        }
        break;
    case DENSE:
        for (int bi = 0; bi < A->m; bi++)
        {
            for (int bj = 0; bj < A->n; bj++)
            {
                if (!A->block[bi * A->n + bj])
                    continue;
                BLOCK *B = A->block[bi * A->n + bj];
                for (int ti = 0; ti < B->m; ti++)
                {
                    for (int tj = 0; tj < B->n; tj++)
                    {
                        if (!B->tile[ti * B->n + tj])
                            continue;
                        fprintf(fp, "%d %d %f\n", bi * A->block_size + ti, bj * A->block_size + tj, 0.0);
                    }
                }
            }
        }

    case CSC:
    for (int i = 0; i < A->real_m; i++)
            {
                for (int j = A->col[i]; j < A->col[i + 1]; j++)
                {
    #if 0
                        fprintf(fp, "%d %d %f\n", i, A->col[j], A->a[j]);
    #else
                    fprintf(fp, "%d %d %f\n", i, A->row[j], 1.0);
    #endif
                }
            }
            break;
#if 0
			for (int i = 0; i < A->m; i++) {
				for (int j = 0; j < A->n; j++) {
					if (!A->block[i * A->n + j]) continue;
					fprintf(fp, "%d %d %f\n", i, j, 0);
				}
			}
#endif
        break;
    default:
        break;
    }

    fclose(fp);

    scatter(A);
}

void information(MATRIX *A)
{
    // compute information
    int A_block = A->num_block;
    int A_tile = A->num_tile;
    double A_mem = A->mem;

    A->num_block = 0;
    A->num_tile = 0;
    A->mem = get_memory() / (1024 * 1024);
    A->time = A->time_L + A->time_U + A->time_C;

    long int nonzero = 0;
    for (int bi = 0; bi < A->m; bi++)
    {
        for (int bj = 0; bj < A->n; bj++)
        {
            BLOCK *B = A->block[bi * A->n + bj];
            if (!B)
                continue;
            A->num_block++;

            for (int ti = 0; ti < B->m; ti++)
            {
                for (int tj = 0; tj < B->n; tj++)
                {
                    TILE *T = B->tile[ti * B->n + tj];
                    if (!T)
                        continue;
                    update_nonz(T);
                    nonzero += T->nonz;
                    A->num_tile++;
                }
            }
        }
    }

    /* print information ***********************************************/
    printf("INFORMATION:\n");

    /* basic information */
    printf("MATRIX: %d x %d, nonz: %d\n", A->real_m, A->real_n, A->real_nonz);
    printf("BLOCK SIZE: %d, TILE_SIZE: %d\n", A->block_size, A->tile_size);

    /* status of block, tile and nonzero element */
    printf("            A        LU       LU/A\n");
    printf("MEM:        %.1e  %.1e  %.1f\n", A_mem, A->mem, A->mem / A_mem);
    printf("BLOCK:      %.1e  %.1e  %.1f\n", (double)A_block, (double)A->num_block, (double)A->num_block / A_block);
    printf("TILE:       %.1e  %.1e  %.1f\n", (double)A_tile, (double)A->num_tile, (double)A->num_tile / A_tile);
    printf("NONZERO:    %.1e  %.1e  %.1f\n", (double)A->real_nonz, (double)nonzero, (double)nonzero / A->real_nonz);
    printf("DENSEITY:   %.1e  %.1e  %.1f\n", (double)A->real_nonz / (A_tile * A->tile_size * A->tile_size),
           (double)nonzero / (A->num_tile * A->tile_size * A->tile_size),
           ((double)nonzero * A_tile) / ((double)A->real_nonz * A->num_tile));

    printf("TIME : %.2f (%.3f, %.3f, %.3f)\n", A->time, A->time_L, A->time_U, A->time_C);
    double flops = 2 * A->tile_size * A->tile_size * A->tile_size;
    double gflops = flops * 1e-9;
    printf("FLOPS: %.2f GFLOPS\n", (A->num_tgemm * gflops) / A->time);

        // 打开文件
    if(A->output_path == NULL){
        return;
    }
    FILE *file = fopen(A->output_path, "a");
    if (file == NULL) {
        perror("无法打开文件");
        exit(EXIT_FAILURE);
    }


    // 写入标题（如果文件为空）
    printf("SAVE LOG : save in %s\n", A->output_path);
    fseek(file, 0, SEEK_END);
    if (ftell(file) == 0) {
      // basic info , MATRIX_name,MATRIX_m,MATRIX_n,real_nonz,block_size,tile_size
        fprintf(file,
                "name,m,n,nonz,block_size,tile_size,");
        // mem info , A_mem,LU_mem,LU/A_mem
        fprintf(file, "A_mem,LU_mem,LU/A_mem,");
        // nonzero info , real_nonz,LU_nonz,LU/A_nonz
        fprintf(file, "real_nonz,LU_nonz,LU/A_nonz,");
        fprintf(file, "reorder, fact,solve,");
        // GFLOPS info , GFLOPS
        fprintf(file, "gflops,");
        // error info, error
        fprintf(file, "error\n");
        }
    
        // 写入数据 
        // basic info , MATRIX_name,MATRIX_m,MATRIX_n,real_nonz,block_size,tile_size
        fprintf(file, "%s,%d,%d,%d,%d,%d,", A->name, A->real_m, A->real_n,
            A->real_nonz, A->block_size, A->tile_size);

        // mem info , A_mem,LU_mem,LU/A_mem
        fprintf(file, "%.5lf,%.5lf,%.1f,", A_mem, A->mem, A->mem / A_mem);

        // nonzero info , real_nonz,LU_nonz,LU/A_nonz
        fprintf(file, "%d,%ld,%lf,", A->real_nonz, nonzero,
            (double)nonzero / A->real_nonz);

        // density info , real_density,LU_density,LU/A_density
        // fprintf(file, "%lf,%lf,%lf,",
        //     (double)A->real_nonz / (A_tile * A->tile_size * A->tile_size),
        //     (double)nonzero / (A->num_tile * A->tile_size * A->tile_size),
        //     ((double)nonzero * A_tile) / ((double)A->real_nonz * A->num_tile));

        fprintf(file, "%.3lf,%.3lf,%.3lf,", A->reorder_time, A->fact_time, A->solve_time);

        // GFLOPS info , GFLOPS
        fprintf(file, "%lf,", (A->num_tgemm * gflops) / A->time);
        fprintf(file, "%lf", A->error);
        fprintf(file, "\n");

            // 关闭文件
        fclose(file);

}
